# -*- coding: utf-8 -*-
# ./fcli shell
# upf wxa_login -h wxa_login.handler -d ./wxa_login
# scp root@testar:/home/dev/code/script/nls/alibabacloud-nls-python-sdk-dev/tests/1.wav ./
# base引入
import json, requests, time, os, re, sys, datetime
from base.webin import *
import base.db as PageDb
import base.mp as MpApi

# nls相关
import threading
import nls
import oss2
# from oss2.credentials import EnvironmentVariableCredentialsProvider

URL="wss://nls-gateway-cn-shanghai.aliyuncs.com/ws/v1"
APPKEY="kuGaAAoN537XcBf4"  #获取Appkey请前往控制台：https://nls-portal.console.aliyun.com/applist
GLOBAL_TIME=0

# openai相关
sys.path.append('/opt/python')
sys.path.append('/opt/lib')
from typing import *
from openai import OpenAI
from openai.types.chat.chat_completion import Choice

'''
chat相关
'''
GLOBAL_TIME=0
def printTime(mark):
    global GLOBAL_TIME
    _time = int(time.time() * 1000)
    if GLOBAL_TIME == 0:
      print("%s===>>:0"% mark)
    else:
      print("%s===>>:%d ms" % (mark, _time - GLOBAL_TIME))   
    GLOBAL_TIME = _time

def contains_punctuation(s):
    punctuation = ".,!?;:，。！？；"  # 您可以在这里添加需要检查的标点符号
    return any(char in punctuation for char in s)

# search 工具的具体实现，这里我们只需要返回参数即可
def search_impl(arguments: Dict[str, Any]) -> Any:
	"""
	在使用 Moonshot AI 提供的 search 工具的场合，只需要原封不动返回 arguments 即可，
	不需要额外的处理逻辑。

	但如果你想使用其他模型，并保留联网搜索的功能，那你只需要修改这里的实现（例如调用搜索
	和获取网页内容等），函数签名不变，依然是 work 的。

	这最大程度保证了兼容性，允许你在不同的模型间切换，并且不需要对代码有破坏性的修改。
	"""
	return arguments

class ChatOpenAI:
    def __init__(self, platform, voiceid):
        self.voiceid = voiceid
        config_map = {
            "deepseek": {
                "model": "deepseek-chat",
                "key": "sk-4013b548f59e47448d928e322e09bab0",
                "base_url": "https://api.deepseek.com"
            },
            "kimi": {
                "model": "moonshot-v1-128k",
                "key": "sk-QuJeZbnZzoMhOs3WYnCRcGeQYA3ySJAXWE1NCb7Ng43iILp1",
                "base_url": "https://api.moonshot.cn/v1"
            },
            "ali": {
                "model": "qwen-plus",
                "key": "sk-c13fb245ab3e4eb19d5c6026597cc70a",
                "base_url": "https://dashscope-finance.aliyuncs.com/api/v1"
            }
        }

        self.config = config_map[platform]

        self.client = OpenAI(
            api_key=self.config["key"],
            base_url=self.config["base_url"],
        )
    
    def get_history_messages(self, input, db):
        now_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
        PREFIX = "你是 星仔，现在是%s，由特星球AI提供的人工智能助手，你更擅长中文的对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力，涉及政治敏感人物等问题的回答。特星球AI 为专有名词，不可翻译成其他语言。" % (now_time)
        self.messages = [
              {"role": "system", "content": PREFIX}
        ]

        # 获取dbs[3]对应id的历史数据
        if db.dbs[3].exists(self.voiceid):
            temp_messages = json.loads( db.dbs[3].get(self.voiceid) )
            messages = []
            for msg in temp_messages:
              if msg.get("content"):
                  messages.append(msg)
            self.messages = messages

            # 重置时间
            self.messages[0] = {"role": "system", "content": PREFIX}

        self.messages.append({"role": "user", "content": input})


    def chat(self, stream=False):
        completion = self.client.chat.completions.create(
            model=self.config["model"],
            messages=self.messages,
            temperature=0.3,
            # tools=[
            #     {
            #         "type": "builtin_function",
            #         "function": {
            #             "name": "$web_search",
            #         },
            #     }
            # ],
            n = 1,
            stream=stream,  # 开启流式输出
        )
        return completion

    def start(self, i, db, text=None):
        voice_index = 0
        response = []
        if text:
          self.get_history_messages(text, db)
        
        messages = [{}]
        temp_content = ""
        completion = self.chat(True)
        for chunk in completion:
          # 通过循环获取每个数据块中所有的 choice，并获取 index 对应的 message 对象
          for choice in chunk.choices:
            index = choice.index
            message = messages[index]
            delta = choice.delta
            role = delta.role
            if role:
                message["role"] = role
            content = delta.content
            if content:
              if not temp_content:
                temp_content = content
              else:
                temp_content = temp_content + content

              if "content" not in message or not message["content"]:
                message["content"] = content
              else:
                message["content"] = message["content"] + content
              
              limit_count = 24
              if voice_index == 0:
                 limit_count = 6
              if (contains_punctuation(content) and len(temp_content) >= limit_count) or not content: #>24个字且包含标点符号则生成语音
                 voiceid = self.voiceid + "_%d" % (voice_index)
                 if len(response) > 0:
                    db.dbs[2].lpush(self.voiceid, json.dumps({
                       "content": temp_content,
                       "voice_index_id": voice_index
                    }))
                    db.dbs[2].expire(self.voiceid, 5 * 60)
                    response.append('http://testars-nls.oss-cn-shenzhen.aliyuncs.com/nls/%s.wav' % (voiceid))
                 else:
                    response.append('http://testars-nls.oss-cn-shenzhen.aliyuncs.com/' + nlsrun(i, db, temp_content, voiceid))
                 print("temp_content==>>", temp_content)
                 temp_content = ""
                 voice_index = voice_index + 1

            # 从这里，我们开始处理 tool_calls
            tool_calls = delta.tool_calls  # <-- 先判断数据块中是否包含 tool_calls
            if tool_calls:
                if "tool_calls" not in message:
                    message["tool_calls"] = []  # <-- 如果包含 tool_calls，我们初始化一个列表来保存这些 tool_calls，注意此时的列表中没有任何元素，长度为 0
                for tool_call in tool_calls:
                    tool_call_index = tool_call.index  # <-- 获取当前 tool_call 的 index 索引
                    if len(message["tool_calls"]) < (
                            tool_call_index + 1):  # <-- 根据 index 索引扩充 tool_calls 列表，以便于我们能通过下标访问到对应的 tool_call
                        message["tool_calls"].extend([{}] * (tool_call_index + 1 - len(message["tool_calls"])))
                    tool_call_object = message["tool_calls"][tool_call_index]  # <-- 根据下标访问对应的 tool_call
                    tool_call_object["index"] = tool_call_index
    
                    # 下面的步骤，是根据数据块中的信息填充每个 tool_call 的 id、type、function 字段
                    # 在 function 字段中，又包括 name 和 arguments 字段，arguments 字段会由每个数据块
                    # 依次补充，如同 delta.content 字段一般。
    
                    tool_call_id = tool_call.id
                    if tool_call_id:
                        tool_call_object["id"] = tool_call_id
                    tool_call_type = tool_call.type
                    if tool_call_type:
                        tool_call_object["type"] = tool_call_type
                    tool_call_function = tool_call.function
                    if tool_call_function:
                        if "function" not in tool_call_object:
                            tool_call_object["function"] = {}
                        tool_call_function_name = tool_call_function.name
                        if tool_call_function_name:
                            tool_call_object["function"]["name"] = tool_call_function_name
                        tool_call_function_arguments = tool_call_function.arguments
                        if tool_call_function_arguments:
                            if "arguments" not in tool_call_object["function"]:
                                tool_call_object["function"]["arguments"] = tool_call_function_arguments
                            else:
                                tool_call_object["function"]["arguments"] = tool_call_object["function"][
                                                                                "arguments"] + tool_call_function_arguments  # <-- 依次补充 function.arguments 字段的值
                    message["tool_calls"][tool_call_index] = tool_call_object
        #
        for index, message in enumerate(messages):
          print("index:", index)
          print("message:", json.dumps(message, ensure_ascii=False))
          if message.get("tool_calls"):
            tool_calls = message['tool_calls']
            # 执行搜索工具调用
            search_results = []
            for tool_call in tool_calls:
                if tool_call['function']['name'] == 'search' or tool_call['function']['name'] == '$web_search':
                    tool_call_arguments = tool_call['function']['arguments']
                    if not tool_call_arguments:
                      tool_call_arguments = {}
                    if isinstance(tool_call_arguments, str):
                      tool_call_arguments = json.loads(tool_call_arguments)
                    if tool_call_arguments.get("query"):
                      query = tool_call_arguments['query']
                      print("query==>>", query)
                      # 假设 search_function 是一个实现搜索的函数
                      search_results = search_impl(query)
                    else:
                      search_results = [tool_call_arguments]
            # # 将搜索结果作为新的上下文传递给模型
            # self.messages.append({
            #     "role": "tool",
            #     "tool_call_id": tool_calls[0]['id'],
            #     "content": search_results
            # })
            # print("self.messages==>>", self.messages)
            time.sleep(10)
            response = self.start2(i, db)
        
        self.messages.extend(messages)
        db.dbs[3].set(self.voiceid, json.dumps( self.messages ))
        db.dbs[3].expire(self.voiceid, 5 * 60)
        return response

    def chat2(self) -> Choice:
      completion = self.client.chat.completions.create(
        model = "moonshot-v1-128k",
        messages = self.messages,
        # max_tokens = 40,
        tools=[
          {
            "type": "builtin_function",  # <-- 使用 builtin_function 声明 $web_search 函数，请在每次请求都完整地带上 tools 声明
            "function": {
              "name": "$web_search",
            },
          }
        ],
        temperature = 0.3,
      )
      # 通过 API 我们获得了 Kimi 大模型给予我们的回复消息（role=assistant）
      usage = completion.usage
      choice = completion.choices[0]
      # msg = choice.message.content
      # print("choice.finish_reason==>>", choice.finish_reason)
      # print(f"chat_prompt_tokens:          {usage.prompt_tokens}")
      # print(f"chat_completion_tokens:      {usage.completion_tokens}")
      # print(f"chat_total_tokens:           {usage.total_tokens}")
      # print("msg==>>", msg)
      return choice

    def start2(self, i, db):
      finish_reason = None
      while finish_reason is None or finish_reason == "tool_calls":
        choice = self.chat2() # self.messages
        finish_reason = choice.finish_reason
        if finish_reason == "tool_calls":
          self.messages.append(choice.message)
          for tool_call in choice.message.tool_calls:
            tool_call_name = tool_call.function.name
            tool_call_arguments = json.loads(
              tool_call.function.arguments)
            if tool_call_name == "$web_search":
              # ===================================================================
              # 我们将联网搜索过程中，由联网搜索结果产生的 Tokens 打印出来
              search_content_total_tokens = tool_call_arguments.get("usage", {}).get("total_tokens")
              # print(f"search_content_total_tokens: {search_content_total_tokens}")
              # ===================================================================
              tool_result = search_impl(tool_call_arguments)
            else:
              tool_result = f"Error: unable to find tool by name '{tool_call_name}'"

            self.messages.append({
              "role": "tool",
              "tool_call_id": tool_call.id,
              "name": tool_call_name,
              "content": json.dumps(tool_result),
            })

      msg = choice.message.content
      return ['http://testars-nls.oss-cn-shenzhen.aliyuncs.com/' + nlsrun(i, db, msg, self.voiceid)]
    
def chatFunc(i, db):
  text = i.args.get('text')
  voiceid = i.args.get("voiceid")
  tts = i.args.get("tts")
  platform = i.args.get('platform', 'kimi')

  if not text:
     return []

  if not voiceid:
    voiceid = "%s" % (i.util.smd5(text))
  
  ai = ChatOpenAI(platform, voiceid)

  response = ai.start(i, db, text) # 向ai提问

  return response


'''
nls相关
'''
#以下代码会根据上述TEXT文本反复进行语音合成
class TtsClass:
  def __init__(self, i, db, result_text, voiceid):
    self.db = db
    self.text = result_text
    self.name = voiceid#i.util.smd5(self.text)
    tid = self.name
    test_file = "/tmp/%s.wav" % (self.name)
    self.__th = threading.Thread(target=self.__test_run)
    self.__id = tid
    self.__test_file = test_file
    self.token = self.get_token()
    self.finish = False

  def get_token(self):
    nls_token_json = json.loads( self.db.dbs[0].hget("aliyunToken", "nls").decode("utf-8") )
    return nls_token_json['token']

  def start(self, text):
    self.__text = text
    self.__f = open(self.__test_file, "wb")
    self.__th.start()
  
  def test_on_metainfo(self, message, *args):
    print("on_metainfo message=>{}".format(message))  

  def test_on_error(self, message, *args):
    print("on_error args=>{}".format(args))

  def test_on_close(self, *args):
    print("on_close: args=>{}".format(args))
    try:
        self.__f.close()
    except Exception as e:
        print("close file failed since:", e)

  def test_on_data(self, data, *args):
    try:
        self.__f.write(data)
    except Exception as e:
        print("write data failed:", e)

  def test_on_completed(self, message, *args):
    print("on_completed:args=>{} message=>{}".format(args, message))

  def __test_run(self):
    print("thread:{} start..".format(self.__id))
    tts = nls.NlsSpeechSynthesizer(url=URL,
                                   token=self.token,
                                   appkey=APPKEY,
                                   on_metainfo=self.test_on_metainfo,
                                   on_data=self.test_on_data,
                                   on_completed=self.test_on_completed,
                                   on_error=self.test_on_error,
                                   on_close=self.test_on_close,
                                   callback_args=[self.__id])
    print("{}: session start".format(self.__id))
    # r = tts.start(self.__text, voice="ailun")
    dictname = {'format': 'wav'}
    r = tts.start(self.__text, voice="zhibei_emo", ex=dictname)
    print("{}: tts done with result:{}".format(self.__id, r))
    self.finish = True

def nlsrun(i, db, result_text, voiceid):
  t = TtsClass(i, db, result_text, voiceid)
  t.start(t.text)
  while not t.finish:
    time.sleep(0.1)
  return save_to_oss(t.name)
    
def save_to_oss(name):
  file_path = "nls/%s.wav" % name
  file_name = "/tmp/%s.wav" % name
  auth = oss2.Auth('LTAI5tE996brnwuKtQKJpqop', 'iXydF3Q5kCPLpwkYpXLIpuQoKBdYEm')
  bucket = oss2.Bucket(auth, 'http://oss-cn-shenzhen.aliyuncs.com', 'testars-nls')
  r = bucket.put_object_from_file(file_path, file_name)
  return file_path

'''
asr相关
'''
#以下代码会根据音频文件内容反复进行一句话识别
class TestSr:
  def __init__(self, db, tid, test_file):
      self.__th = threading.Thread(target=self.__test_run)
      self.__id = tid
      self.__test_file = "/tmp/%s.wav" % test_file
      self.db = db
      self.token = self.get_token()
      self.finish = False
      self.result = ""
    
  def get_token(self):
      nls_token_json = json.loads( self.db.dbs[0].hget("aliyunToken", "nls").decode("utf-8") )
      return nls_token_json['token']
  
  def loadfile(self, filename):
      with open(filename, "rb") as f:
          self.__data = f.read()
  
  def start(self):
      self.loadfile(self.__test_file)
      self.__th.start()

  def test_on_start(self, message, *args):
      print("test_on_start:{}".format(message))

  def test_on_error(self, message, *args):
      print("on_error args=>{}".format(args))

  def test_on_close(self, *args):
      print("on_close: args=>{}".format(args))

  def test_on_result_chg(self, message, *args):
      print("test_on_chg:{}".format(message))

  def test_on_completed(self, message, *args):
      print("on_completed:args=>{} message=>{}".format(args, message))
      message = json.loads(message)
      print("message==>>", message)
      if message.get("payload") and message["payload"].get("result"):
        self.result = message["payload"]["result"]

  def __test_run(self):
      print("thread:{} start..".format(self.__id))
      
      sr = nls.NlsSpeechRecognizer(
                  url=URL,
                  token=self.token,
                  appkey=APPKEY,
                  on_start=self.test_on_start,
                  on_result_changed=self.test_on_result_chg,
                  on_completed=self.test_on_completed,
                  on_error=self.test_on_error,
                  on_close=self.test_on_close,
                  callback_args=[self.__id]
              )

      print("{}: session start".format(self.__id))
      r = sr.start(aformat="pcm", ex={"hello":123})
          
      self.__slices = zip(*(iter(self.__data),) * 640)
      for i in self.__slices:
          sr.send_audio(bytes(i))
          time.sleep(0.01)

      r = sr.stop()
      print("{}: sr stopped:{}".format(self.__id, r))
      self.finish = True

def download_from_oss(name):
  file_path = "asr/%s.wav" % name
  file_name = "/tmp/%s.wav" % name
  print("file_path==>>", file_path)
  auth = oss2.Auth('LTAI5tE996brnwuKtQKJpqop', 'iXydF3Q5kCPLpwkYpXLIpuQoKBdYEm')
  bucket = oss2.Bucket(auth, 'http://oss-cn-shenzhen.aliyuncs.com', 'testars-nls')
  bucket.get_object_to_file(file_path, file_name)
  return file_name

def asrrun(db, file_name):
  download_from_oss(file_name)
  name = "thread_1"
  t = TestSr(db, name, file_name)
  t.start()
  while not t.finish:
    time.sleep(0.1)
  print("t.result==>>", t.result)
  return t.result

def handler(environ, start_response):
  i = Webin(environ)
  printTime("函数计算开始")
  text = i.args.get('text')
  tts = i.args.get('tts')
  asr = i.args.get('asr')
  if not (text or tts or asr):
    status = '200 OK'
    response_headers = [('Content-type', 'text/plain')]
    # start_response(status, response_headers)
    return json.dumps(i.util.return_fail("no text"))
  
  nls.enableTrace(True)
  db = PageDb.PageDb("system")
  
  if asr:
    printTime("asr开始")
    i.args['text'] = asrrun(db, asr)
    printTime("asr结束")
    response = chatFunc(i, db)
  else:
    response = chatFunc(i, db)

  printTime("函数计算返回")
  status = '200 OK'
  response_headers = [('Content-type', 'text/plain')]
  # start_response(status, response_headers)
  return i.util.return_success(response)
